import torch.nn as nn


class CNN_3D(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv11 = nn.Conv3d(1, 16, (4, 9, 9), stride=(1, 2, 1))
        self.conv11_bn = nn.BatchNorm3d(16)
        self.conv11_activation = nn.PReLU()
        self.conv12 = nn.Conv3d(16, 16, (4, 9, 9), stride=(1, 1, 1))
        self.conv12_bn = nn.BatchNorm3d(16)
        self.conv12_activation = nn.PReLU()
        self.conv21 = nn.Conv3d(16, 32, (3, 7, 7), stride=(1, 1, 1))
        self.conv21_bn = nn.BatchNorm3d(32)
        self.conv21_activation = nn.PReLU()
        self.conv22 = nn.Conv3d(32, 32, (3, 7, 7), stride=(1, 1, 1))
        self.conv22_bn = nn.BatchNorm3d(32)
        self.conv22_activation = nn.PReLU()
        self.conv31 = nn.Conv3d(32, 64, (3, 5, 5), stride=(1, 1, 1))
        self.conv31_bn = nn.BatchNorm3d(64)
        self.conv31_activation = nn.PReLU()
        self.conv32 = nn.Conv3d(64, 64, (3, 5, 5), stride=(1, 1, 1))
        self.conv32_bn = nn.BatchNorm3d(64)
        self.conv32_activation = nn.PReLU()
        self.conv41 = nn.Conv3d(64, 128, (3, 3, 3), stride=(1, 1, 1))
        self.conv41_bn = nn.BatchNorm3d(128)
        self.conv41_activation = nn.PReLU()

    def forward(self, x):
        x = self.conv11(x)
        x = self.conv11_bn(x)
        x = self.conv11_activation(x)
        x = self.conv12(x)
        x = self.conv12_bn(x)
        x = self.conv12_activation(x)
        x = self.conv21(x)
        x = self.conv21_bn(x)
        x = self.conv21_activation(x)
        x = self.conv22(x)
        x = self.conv22_bn(x)
        x = self.conv22_activation(x)
        x = self.conv31(x)
        x = self.conv31_bn(x)
        x = self.conv31_activation(x)
        x = self.conv32(x)
        x = self.conv32_bn(x)
        x = self.conv32_activation(x)
        x = self.conv41(x)
        x = self.conv41_bn(x)
        x = self.conv41_activation(x)
        return x
